{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../scripts/')\n",
    "from dp_policy_agent import *\n",
    "import random, copy "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StateInfo: #q4stateinfo\n",
    "    def __init__(self, action_num, epsilon=0.3):\n",
    "        self.q = np.zeros(action_num)\n",
    "        self.epsilon = epsilon\n",
    "        \n",
    "    def greedy(self):\n",
    "        return np.argmax(self.q)\n",
    "    \n",
    "    def epsilon_greedy(self, epsilon): \n",
    "        if random.random() < epsilon:\n",
    "            return random.choice(range(len(self.q)))\n",
    "        else:\n",
    "            return self.greedy()\n",
    "    \n",
    "    def pi(self):\n",
    "        return self.epsilon_greedy(self.epsilon) \n",
    "    \n",
    "    def max_q(self):              #追加\n",
    "        return max(self.q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SarsaAgent(PuddleIgnoreAgent): #名前を変えましょう ###sarsa1 \n",
    "    def __init__(self, time_interval, particle_pose, envmap, puddle_coef=100, alpha=0.5, \\\n",
    "                motion_noise_stds={\"nn\":0.19, \"no\":0.001, \"on\":0.13, \"oo\":0.2}, widths=np.array([0.2, 0.2, math.pi/18]).T, \\\n",
    "                lowerleft=np.array([-4, -4]).T, upperright=np.array([4, 4]).T):   #alpha追加\n",
    " \n",
    "        super().__init__(time_interval, particle_pose, envmap, None, puddle_coef, motion_noise_stds)\n",
    "        \n",
    "        ###DynamicProgrammingから持ってくる###\n",
    "        self.pose_min = np.r_[lowerleft, 0]\n",
    "        self.pose_max = np.r_[upperright, math.pi*2]\n",
    "        self.widths = widths\n",
    "        self.index_nums = ((self.pose_max - self.pose_min)/self.widths).astype(int)\n",
    "        nx, ny, nt = self.index_nums\n",
    "        self.indexes = list(itertools.product(range(nx), range(ny), range(nt)))\n",
    "        \n",
    "        ###PuddleIgnorePolicyの方策と価値関数の読み込み###\n",
    "        self.actions, self.ss = self.set_action_value_function()\n",
    "        \n",
    "        ###強化学習用変数### #追加\n",
    "        self.alpha = alpha\n",
    "        self.s, self.a = None, None\n",
    "        self.update_end = False\n",
    "        \n",
    "        self.step = 0\n",
    "\n",
    "    def set_action_value_function(self): \n",
    "        policy = np.zeros(np.r_[self.index_nums,2])\n",
    "        for line in open(\"puddle_ignore_policy.txt\", \"r\"): \n",
    "            d = line.split()\n",
    "            policy[int(d[0]), int(d[1]), int(d[2])] = [float(d[3]), float(d[4])]\n",
    "        \n",
    "        actions = list(set([tuple(policy[i]) for i in self.indexes])) \n",
    "        action_num = len(actions)\n",
    "        \n",
    "        ss = {}\n",
    "        for line in open(\"puddle_ignore_values.txt\", \"r\"): \n",
    "            d = line.split()\n",
    "            index, value = (int(d[0]), int(d[1]), int(d[2])), float(d[3]) \n",
    "            ss[index] = StateInfo(action_num) \n",
    "            \n",
    "            for i, a in enumerate(actions): \n",
    "                ss[index].q[i] = value if tuple(policy[index]) == a else value - 0.1\n",
    "                ss[index].q[i] *= 10 #大きくしておく\n",
    "                \n",
    "        return actions, ss\n",
    "    \n",
    "    def policy(self, pose): ###q4policy\n",
    "        index = np.floor((pose - self.pose_min)/self.widths).astype(int)\n",
    "        \n",
    "        index[2] = (index[2] + self.index_nums[2]*1000)%self.index_nums[2]\n",
    "        for i in [0,1]:\n",
    "            if index[i] < 0: index[i] = 0\n",
    "            elif index[i] >= self.index_nums[i]: index[i] = self.index_nums[i] - 1\n",
    "                \n",
    "        s = tuple(index)          #Q学習の式で使う記号に変更\n",
    "        a = self.ss[s].pi()\n",
    "        return s, a                 #行動は値でなくインデックスで返す\n",
    "            \n",
    "    def decision(self, observation=None):###q4decision\n",
    "        if self.step%100000 == 0:\n",
    "            with open(\"tmp/sarsa0_policy\"+ str(self.step)+  \".txt\", \"w\") as f:  ###\n",
    "                for index in self.indexes:\n",
    "                    p = self.actions[self.ss[index].greedy()]\n",
    "                    f.write(\"{} {} {} {} {}\\n\".format(index[0], index[1], index[2], p[0], p[1]))  \n",
    "                f.flush()\n",
    "                    \n",
    "            with open(\"tmp/sarsa0_value\"+ str(self.step)+  \".txt\", \"w\") as f:  ###\n",
    "                for index in self.indexes:\n",
    "                    f.write(\"{} {} {} {}\\n\".format(index[0], index[1], index[2], self.ss[index].max_q()))  \n",
    "                f.flush()\n",
    "                \n",
    "        ###終了処理###\n",
    "        if self.update_end:\n",
    "            return 0.0, 0.0\n",
    "        if self.in_goal:\n",
    "            self.update_end = True #ゴールに入った後も一回だけ更新があるので即終了はしない\n",
    "        \n",
    "        ###カルマンフィルタの実行###\n",
    "        self.kf.motion_update(self.prev_nu, self.prev_omega, self.time_interval)\n",
    "        self.kf.observation_update(observation)\n",
    "        \n",
    "        ###行動決定と報酬の処理###\n",
    "        s_, a_ = self.policy(self.kf.belief.mean)\n",
    "        r = self.time_interval*self.reward_per_sec()\n",
    "        self.total_reward += r\n",
    "        \n",
    "        ###Q学習と現在の状態と行動の保存###\n",
    "        self.q_update(r, s_, a_)      #a_も必要になる\n",
    "        self.s, self.a = s_, a_\n",
    "\n",
    "        ###出力###\n",
    "        self.prev_nu, self.prev_omega = self.actions[a_]\n",
    "        \n",
    "        self.step += 1\n",
    "        return self.actions[a_]\n",
    "\n",
    "    def q_update(self, r, s_, a_):###sarsa1\n",
    "        if self.s == None:\n",
    "            return\n",
    "        \n",
    "        q = self.ss[self.s].q[self.a]\n",
    "        q_ = self.final_value if self.in_goal else self.ss[s_].q[a_]     #max_qからQ(s_,a_)に書き換え\n",
    "        self.ss[self.s].q[self.a] = (1.0 - self.alpha)*q + self.alpha*(r + q_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WarpRobot(Robot): ###q5warprobot\n",
    "    def __init__(self, *args, **kwargs): \n",
    "        super().__init__(*args, **kwargs)\n",
    "        \n",
    "        self.init_agent = copy.deepcopy(self.agent) #エージェントのコピーを残しておく\n",
    "        \n",
    "    def choose_pose(self):  #初期位置をランダムに決めるメソッド（雑）\n",
    "        xy = random.random()*8-4\n",
    "        t = random.random()*2*math.pi\n",
    "        return np.array([4, xy, t]).T if random.random() > 0.5 else np.array([xy, 4, t]).T\n",
    "        \n",
    "    def reset(self):\n",
    "        #ssだけ残してエージェントを初期化\n",
    "        tmp = self.agent.ss\n",
    "        t = self.agent.step\n",
    "        self.agent = copy.deepcopy(self.init_agent)\n",
    "        self.agent.ss  = tmp\n",
    "        self.agent.step = t\n",
    "        \n",
    "        #初期位置をセット（ロボット、カルマンフィルタ）\n",
    "        self.pose = self.choose_pose()\n",
    "        self.agent.kf.belief = multivariate_normal(mean=self.pose, cov=np.diag([1e-10, 1e-10, 1e-10]))\n",
    "        \n",
    "        #軌跡の黒い線が残らないように消す\n",
    "        self.poses = []\n",
    "    \n",
    "    def one_step(self, time_interval):\n",
    "        if self.agent.update_end:\n",
    "            with open(\"log.txt\", \"a\") as f:\n",
    "                f.write(\"{}\\n\".format(self.agent.total_reward + self.agent.final_value))\n",
    "            self.reset()\n",
    "            return\n",
    "        \n",
    "        super().one_step(time_interval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "__init__() takes from 4 to 5 positional arguments but 7 were given",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-bbc8fba6c39f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0;31m###ロボットを1台登場させる###\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0minit_pose\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m3\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mT\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m     \u001b[0msa\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mSarsaAgent\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtime_interval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minit_pose\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mm\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m#エージェント変更\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m     r = WarpRobot(init_pose, sensor=Camera(m, distance_bias_rate_stddev=0, direction_bias_stddev=0), \n\u001b[1;32m     24\u001b[0m               agent=sa, color=\"red\", bias_rate_stds=(0,0))\n",
      "\u001b[0;32m<ipython-input-3-7c393b2bfab7>\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, time_interval, particle_pose, envmap, puddle_coef, alpha, motion_noise_stds, widths, lowerleft, upperright)\u001b[0m\n\u001b[1;32m      4\u001b[0m                 lowerleft=np.array([-4, -4]).T, upperright=np.array([4, 4]).T):   #alpha追加\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m         \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtime_interval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparticle_pose\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0menvmap\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpuddle_coef\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmotion_noise_stds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0;31m###DynamicProgrammingから持ってくる###\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: __init__() takes from 4 to 5 positional arguments but 7 were given"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__': \n",
    "    time_interval = 0.1\n",
    "    world = PuddleWorld(400000000, time_interval) #時間をほぼ無限に\n",
    "\n",
    "    m = Map()\n",
    "    m.append_landmark(Landmark(-4,2))\n",
    "    m.append_landmark(Landmark(2,-3))\n",
    "    m.append_landmark(Landmark(4,4))\n",
    "    m.append_landmark(Landmark(-4,-4))\n",
    "    world.append(m)\n",
    "    \n",
    "    ###ゴールの追加###\n",
    "    goal = Goal(-3,-3)\n",
    "    world.append(goal)\n",
    "    \n",
    "    ###水たまりの追加###\n",
    "    world.append(Puddle((-2, 0), (0, 2), 0.1)) \n",
    "    world.append(Puddle((-0.5, -2), (2.5, 1), 0.1))\n",
    "\n",
    "    ###ロボットを1台登場させる###\n",
    "    init_pose = np.array([3, 3, 0]).T\n",
    "    sa = SarsaAgent(time_interval, init_pose, m) #エージェント変更\n",
    "    r = WarpRobot(init_pose, sensor=Camera(m, distance_bias_rate_stddev=0, direction_bias_stddev=0), \n",
    "              agent=sa, color=\"red\", bias_rate_stds=(0,0))\n",
    "\n",
    "    world.append(r)\n",
    "    \n",
    "    world.draw()\n",
    "    #r.one_step(0.1) #デバッグ時"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = np.zeros(sa.index_nums[0:2])\n",
    "for x in range(sa.index_nums[0]):\n",
    "    for y in range(sa.index_nums[1]):\n",
    "        a = sa.ss[(x,y,22)].greedy()\n",
    "        p[x,y] = sa.actions[a][0] + sa.actions[a][1]\n",
    "        \n",
    "import seaborn as sns   \n",
    "sns.heatmap(np.rot90(p), square=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = np.zeros(r.agent.index_nums[0:2])\n",
    "for x in range(r.agent.index_nums[0]):\n",
    "    for y in range(r.agent.index_nums[1]):\n",
    "        v[x,y] = r.agent.ss[(x,y,18)].max_q()\n",
    "        \n",
    "import seaborn as sns   \n",
    "sns.heatmap(np.rot90(v), square=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"sarsa_policy0.txt\", \"w\") as f:  ###\n",
    "    for index in sa.indexes:\n",
    "        p = sa.actions[sa.ss[index].greedy()]\n",
    "        f.write(\"{} {} {} {} {}\\n\".format(index[0], index[1], index[2], p[0], p[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
